{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "bert.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "j0a4mTk9o1Qg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Modified source from https://colab.research.google.com/github/google-research/bert/blob/master/predicting_movie_reviews_with_bert_on_tf_hub.ipynb\n",
        "\n",
        "# Copyright 2019 Google Inc.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCpvgG0vwXAZ",
        "colab_type": "text"
      },
      "source": [
        "#Predicting Movie Review Sentiment with BERT on TF Hub",
        "\n",
        "_WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version`_\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xiYrZKaHwV81",
        "colab_type": "text"
      },
      "source": [
        "If you’ve been following Natural Language Processing over the past year, you’ve probably heard of BERT: Bidirectional Encoder Representations from Transformers. It’s a neural network architecture designed by Google researchers that’s totally transformed what’s state-of-the-art for NLP tasks, like text classification, translation, summarization, and question answering.\n",
        "\n",
        "Now that BERT's been added to [TF Hub](https://www.tensorflow.org/hub) as a loadable module, it's easy(ish) to add into existing TensorFlow text pipelines. In an existing pipeline, BERT can replace text embedding layers like ELMO and GloVE. Alternatively, [finetuning](http://wiki.fast.ai/index.php/Fine_tuning) BERT can provide both an accuracy boost and faster training time in many cases.\n",
        "\n",
        "Here, we'll train a model to predict whether an IMDB movie review is positive or negative using BERT in TensorFlow with tf hub. Some code was adapted from [this colab notebook](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb). Let's get started!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "chM4UttbMIqq",
        "colab_type": "text"
      },
      "source": [
        "First, we'll install the required packages:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jviywGyWyKsA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install bert-tensorflow==1.0.* tensorflow-gpu==1.13.* scikit-learn==0.21.* pandas==0.24.* tensorflow-hub==0.6.* boto3==1.*"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hsZvic2YxnTz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from datetime import datetime\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import bert\n",
        "from bert import run_classifier\n",
        "from bert import optimization\n",
        "from bert import tokenization"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KVB3eOcjxxm1",
        "colab_type": "text"
      },
      "source": [
        "Below, we'll set an output location to store our model output, checkpoints, and export in a local directory. Note: if you're running on Google Colab, local directories don't persist after the session ends."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "US_EAnICvP7f",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "OUTPUT_DIR = \"bert\"\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pmFYvkylMwXn",
        "colab_type": "text"
      },
      "source": [
        "#Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MC_w8SRqN0fr",
        "colab_type": "text"
      },
      "source": [
        "First, let's download the dataset, hosted by Stanford. The code below, which downloads, extracts, and imports the IMDB Large Movie Review Dataset, is borrowed from [this TensorFlow tutorial](https://www.tensorflow.org/hub/tutorials/text_classification_with_tf_hub)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fom_ff20gyy6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tensorflow import keras\n",
        "import os\n",
        "import re\n",
        "\n",
        "# Load all files from a directory in a DataFrame.\n",
        "def load_directory_data(directory):\n",
        "  data = {}\n",
        "  data[\"sentence\"] = []\n",
        "  data[\"sentiment\"] = []\n",
        "  for file_path in os.listdir(directory):\n",
        "    with tf.gfile.GFile(os.path.join(directory, file_path), \"r\") as f:\n",
        "      data[\"sentence\"].append(f.read())\n",
        "      data[\"sentiment\"].append(re.match(\"\\d+_(\\d+)\\.txt\", file_path).group(1))\n",
        "  return pd.DataFrame.from_dict(data)\n",
        "\n",
        "# Merge positive and negative examples, add a polarity column and shuffle.\n",
        "def load_dataset(directory):\n",
        "  pos_df = load_directory_data(os.path.join(directory, \"pos\"))\n",
        "  neg_df = load_directory_data(os.path.join(directory, \"neg\"))\n",
        "  pos_df[\"polarity\"] = 1\n",
        "  neg_df[\"polarity\"] = 0\n",
        "  return pd.concat([pos_df, neg_df]).sample(frac=1).reset_index(drop=True)\n",
        "\n",
        "# Download and process the dataset files.\n",
        "def download_and_load_datasets(force_download=False):\n",
        "  dataset = tf.keras.utils.get_file(\n",
        "      fname=\"aclImdb.tar.gz\", \n",
        "      origin=\"http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\", \n",
        "      extract=True)\n",
        "  \n",
        "  train_df = load_dataset(os.path.join(os.path.dirname(dataset), \n",
        "                                       \"aclImdb\", \"train\"))\n",
        "  test_df = load_dataset(os.path.join(os.path.dirname(dataset), \n",
        "                                      \"aclImdb\", \"test\"))\n",
        "  \n",
        "  return train_df, test_df\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2abfwdn-g135",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train, test = download_and_load_datasets()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XA8WHJgzhIZf",
        "colab_type": "text"
      },
      "source": [
        "To keep training fast, we'll take a sample of 5000 train and test examples, respectively."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lw_F488eixTV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train = train.sample(5000)\n",
        "test = test.sample(5000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "prRQM8pDi8xI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train.columns"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sfRnHSz3iSXz",
        "colab_type": "text"
      },
      "source": [
        "For us, our input data is the 'sentence' column and our label is the 'polarity' column (0, 1 for negative and positive, respecitvely)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IuMOGwFui4it",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "DATA_COLUMN = 'sentence'\n",
        "LABEL_COLUMN = 'polarity'\n",
        "# label_list is the list of labels, i.e. True, False or 0, 1 or 'dog', 'cat'\n",
        "label_list = [0, 1]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V399W0rqNJ-Z",
        "colab_type": "text"
      },
      "source": [
        "#Data Preprocessing\n",
        "We'll need to transform our data into a format BERT understands. This involves two steps. First, we create  `InputExample`'s using the constructor provided in the BERT library.\n",
        "\n",
        "- `text_a` is the text we want to classify, which in this case, is the `Request` field in our Dataframe. \n",
        "- `text_b` is used if we're training a model to understand the relationship between sentences (i.e. is `text_b` a translation of `text_a`? Is `text_b` an answer to the question asked by `text_a`?). This doesn't apply to our task, so we can leave `text_b` blank.\n",
        "- `label` is the label for our example, i.e. True, False"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9gEt5SmM6i6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Use the InputExample class from BERT's run_classifier code to create examples from the data\n",
        "train_InputExamples = train.apply(lambda x: bert.run_classifier.InputExample(guid=None, # Globally unique ID for bookkeeping, unused in this example\n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN]), axis = 1)\n",
        "\n",
        "test_InputExamples = test.apply(lambda x: bert.run_classifier.InputExample(guid=None, \n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN]), axis = 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SCZWZtKxObjh",
        "colab_type": "text"
      },
      "source": [
        "Next, we need to preprocess our data so that it matches the data BERT was trained on. For this, we'll need to do a couple of things (but don't worry--this is also included in the Python library):\n",
        "\n",
        "\n",
        "1. Lowercase our text (if we're using a BERT lowercase model)\n",
        "2. Tokenize it (i.e. \"sally says hi\" -> [\"sally\", \"says\", \"hi\"])\n",
        "3. Break words into WordPieces (i.e. \"calling\" -> [\"call\", \"##ing\"])\n",
        "4. Map our words to indexes using a vocab file that BERT provides\n",
        "5. Add special \"CLS\" and \"SEP\" tokens (see the [readme](https://github.com/google-research/bert))\n",
        "6. Append \"index\" and \"segment\" tokens to each input (see the [BERT paper](https://arxiv.org/pdf/1810.04805.pdf))\n",
        "\n",
        "Happily, we don't have to worry about most of these details.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qMWiDtpyQSoU",
        "colab_type": "text"
      },
      "source": [
        "To start, we'll need to load a vocabulary file and lowercasing information directly from the BERT tf hub module:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IhJSe0QHNG7U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# This is a path to an uncased (all lowercase) version of BERT\n",
        "BERT_MODEL_HUB = \"https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1\"\n",
        "\n",
        "def create_tokenizer_from_hub_module():\n",
        "  \"\"\"Get the vocab file and casing info from the Hub module.\"\"\"\n",
        "  with tf.Graph().as_default():\n",
        "    bert_module = hub.Module(BERT_MODEL_HUB)\n",
        "    tokenization_info = bert_module(signature=\"tokenization_info\", as_dict=True)\n",
        "    with tf.Session() as sess:\n",
        "      vocab_file, do_lower_case = sess.run([tokenization_info[\"vocab_file\"],\n",
        "                                            tokenization_info[\"do_lower_case\"]])\n",
        "      \n",
        "  return bert.tokenization.FullTokenizer(\n",
        "      vocab_file=vocab_file, do_lower_case=do_lower_case)\n",
        "\n",
        "tokenizer = create_tokenizer_from_hub_module()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4oFkhpZBDKm",
        "colab_type": "text"
      },
      "source": [
        "Great--we just learned that the BERT model we're using expects lowercase data (that's what stored in tokenization_info[\"do_lower_case\"]) and we also loaded BERT's vocab file. We also created a tokenizer, which breaks words into word pieces:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dsBo6RCtQmwx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tokenizer.tokenize(\"This here's an example of using the BERT tokenizer\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0OEzfFIt6GIc",
        "colab_type": "text"
      },
      "source": [
        "Using our tokenizer, we'll call `run_classifier.convert_examples_to_features` on our InputExamples to convert them into features BERT understands."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LL5W8gEGRTAf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# We'll set sequences to be at most 128 tokens long.\n",
        "MAX_SEQ_LENGTH = 128\n",
        "# Convert our train and test features to InputFeatures that BERT understands.\n",
        "train_features = bert.run_classifier.convert_examples_to_features(train_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)\n",
        "test_features = bert.run_classifier.convert_examples_to_features(test_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ccp5trMwRtmr",
        "colab_type": "text"
      },
      "source": [
        "#Creating a model\n",
        "\n",
        "Now that we've prepared our data, let's focus on building a model. `create_model` does just this below. First, it loads the BERT tf hub module again (this time to extract the computation graph). Next, it creates a single new layer that will be trained to adapt BERT to our sentiment task (i.e. classifying whether a movie review is positive or negative). This strategy of using a mostly trained model is called [fine-tuning](http://wiki.fast.ai/index.php/Fine_tuning)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6o2a5ZIvRcJq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def create_model(is_predicting, input_ids, input_mask, segment_ids, labels,\n",
        "                 num_labels):\n",
        "  \"\"\"Creates a classification model.\"\"\"\n",
        "\n",
        "  bert_module = hub.Module(\n",
        "      BERT_MODEL_HUB,\n",
        "      trainable=True)\n",
        "  bert_inputs = dict(\n",
        "      input_ids=input_ids,\n",
        "      input_mask=input_mask,\n",
        "      segment_ids=segment_ids)\n",
        "  bert_outputs = bert_module(\n",
        "      inputs=bert_inputs,\n",
        "      signature=\"tokens\",\n",
        "      as_dict=True)\n",
        "\n",
        "  # Use \"pooled_output\" for classification tasks on an entire sentence.\n",
        "  # Use \"sequence_outputs\" for token-level output.\n",
        "  output_layer = bert_outputs[\"pooled_output\"]\n",
        "\n",
        "  hidden_size = output_layer.shape[-1].value\n",
        "\n",
        "  # Create our own layer to tune for politeness data.\n",
        "  output_weights = tf.get_variable(\n",
        "      \"output_weights\", [num_labels, hidden_size],\n",
        "      initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
        "\n",
        "  output_bias = tf.get_variable(\n",
        "      \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\n",
        "\n",
        "  with tf.variable_scope(\"loss\"):\n",
        "\n",
        "    # Dropout helps prevent overfitting\n",
        "    output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)\n",
        "\n",
        "    logits = tf.matmul(output_layer, output_weights, transpose_b=True)\n",
        "    logits = tf.nn.bias_add(logits, output_bias)\n",
        "    log_probs = tf.nn.log_softmax(logits, axis=-1)\n",
        "\n",
        "    # Convert labels into one-hot encoding\n",
        "    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\n",
        "\n",
        "    predicted_labels = tf.squeeze(tf.argmax(log_probs, axis=-1, output_type=tf.int32))\n",
        "    # If we're predicting, we want predicted labels and the probabiltiies.\n",
        "    if is_predicting:\n",
        "      return (predicted_labels, log_probs)\n",
        "\n",
        "    # If we're train/eval, compute loss between predicted and actual label\n",
        "    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\n",
        "    loss = tf.reduce_mean(per_example_loss)\n",
        "    return (loss, predicted_labels, log_probs)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qpE0ZIDOCQzE",
        "colab_type": "text"
      },
      "source": [
        "Next we'll wrap our model function in a `model_fn_builder` function that adapts our model to work for training, evaluation, and prediction."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FnH-AnOQ9KKW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# model_fn_builder actually creates our model function\n",
        "# using the passed parameters for num_labels, learning_rate, etc.\n",
        "def model_fn_builder(num_labels, learning_rate, num_train_steps,\n",
        "                     num_warmup_steps):\n",
        "  \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n",
        "  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
        "    \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n",
        "\n",
        "    input_ids = features[\"input_ids\"]\n",
        "    input_mask = features[\"input_mask\"]\n",
        "    segment_ids = features[\"segment_ids\"]\n",
        "    label_ids = features[\"label_ids\"]\n",
        "\n",
        "    is_predicting = (mode == tf.estimator.ModeKeys.PREDICT)\n",
        "    \n",
        "    # TRAIN and EVAL\n",
        "    if not is_predicting:\n",
        "\n",
        "      (loss, predicted_labels, log_probs) = create_model(\n",
        "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\n",
        "\n",
        "      train_op = bert.optimization.create_optimizer(\n",
        "          loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu=False)\n",
        "\n",
        "      # Calculate evaluation metrics. \n",
        "      def metric_fn(label_ids, predicted_labels):\n",
        "        accuracy = tf.metrics.accuracy(label_ids, predicted_labels)\n",
        "        f1_score = tf.contrib.metrics.f1_score(\n",
        "            label_ids,\n",
        "            predicted_labels)\n",
        "        auc = tf.metrics.auc(\n",
        "            label_ids,\n",
        "            predicted_labels)\n",
        "        recall = tf.metrics.recall(\n",
        "            label_ids,\n",
        "            predicted_labels)\n",
        "        precision = tf.metrics.precision(\n",
        "            label_ids,\n",
        "            predicted_labels) \n",
        "        true_pos = tf.metrics.true_positives(\n",
        "            label_ids,\n",
        "            predicted_labels)\n",
        "        true_neg = tf.metrics.true_negatives(\n",
        "            label_ids,\n",
        "            predicted_labels)   \n",
        "        false_pos = tf.metrics.false_positives(\n",
        "            label_ids,\n",
        "            predicted_labels)  \n",
        "        false_neg = tf.metrics.false_negatives(\n",
        "            label_ids,\n",
        "            predicted_labels)\n",
        "        return {\n",
        "            \"eval_accuracy\": accuracy,\n",
        "            \"f1_score\": f1_score,\n",
        "            \"auc\": auc,\n",
        "            \"precision\": precision,\n",
        "            \"recall\": recall,\n",
        "            \"true_positives\": true_pos,\n",
        "            \"true_negatives\": true_neg,\n",
        "            \"false_positives\": false_pos,\n",
        "            \"false_negatives\": false_neg\n",
        "        }\n",
        "\n",
        "      eval_metrics = metric_fn(label_ids, predicted_labels)\n",
        "\n",
        "      if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "        return tf.estimator.EstimatorSpec(mode=mode,\n",
        "          loss=loss,\n",
        "          train_op=train_op)\n",
        "      else:\n",
        "          return tf.estimator.EstimatorSpec(mode=mode,\n",
        "            loss=loss,\n",
        "            eval_metric_ops=eval_metrics)\n",
        "    else:\n",
        "      (predicted_labels, log_probs) = create_model(\n",
        "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\n",
        "\n",
        "      predictions = {\n",
        "          'probabilities': log_probs,\n",
        "          'labels': predicted_labels\n",
        "      }\n",
        "      return tf.estimator.EstimatorSpec(mode, predictions=predictions)\n",
        "\n",
        "  # Return the actual model function in the closure\n",
        "  return model_fn\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OjwJ4bTeWXD8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Compute train and warmup steps from batch size\n",
        "# These hyperparameters are copied from this colab notebook (https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb)\n",
        "BATCH_SIZE = 32\n",
        "LEARNING_RATE = 2e-5\n",
        "NUM_TRAIN_EPOCHS = 3.0\n",
        "# Warmup is a period of time where hte learning rate \n",
        "# is small and gradually increases--usually helps training.\n",
        "WARMUP_PROPORTION = 0.1\n",
        "# Model configs\n",
        "SAVE_CHECKPOINTS_STEPS = 500\n",
        "SAVE_SUMMARY_STEPS = 100"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "emHf9GhfWBZ_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Compute # train and warmup steps from batch size\n",
        "num_train_steps = int(len(train_features) / BATCH_SIZE * NUM_TRAIN_EPOCHS)\n",
        "num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oEJldMr3WYZa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Specify outpit directory and number of checkpoint steps to save\n",
        "run_config = tf.estimator.RunConfig(\n",
        "    model_dir=OUTPUT_DIR,\n",
        "    save_summary_steps=SAVE_SUMMARY_STEPS,\n",
        "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q_WebpS1X97v",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_fn = model_fn_builder(\n",
        "  num_labels=len(label_list),\n",
        "  learning_rate=LEARNING_RATE,\n",
        "  num_train_steps=num_train_steps,\n",
        "  num_warmup_steps=num_warmup_steps)\n",
        "\n",
        "estimator = tf.estimator.Estimator(\n",
        "  model_fn=model_fn,\n",
        "  config=run_config,\n",
        "  params={\"batch_size\": BATCH_SIZE})\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NOO3RfG1DYLo",
        "colab_type": "text"
      },
      "source": [
        "Next we create an input builder function that takes our training feature set (`train_features`) and produces a generator. This is a pretty standard design pattern for working with TensorFlow [Estimators](https://www.tensorflow.org/guide/estimators)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Pv2bAlOX_-K",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create an input function for training. drop_remainder = True for using TPUs.\n",
        "train_input_fn = bert.run_classifier.input_fn_builder(\n",
        "    features=train_features,\n",
        "    seq_length=MAX_SEQ_LENGTH,\n",
        "    is_training=True,\n",
        "    drop_remainder=False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t6Nukby2EB6-",
        "colab_type": "text"
      },
      "source": [
        "Now we train our model! For me, using a Colab notebook running on Google's GPUs, training time is typically 8-14 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nucD4gluYJmK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(f'Beginning Training!')\n",
        "current_time = datetime.now()\n",
        "estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)\n",
        "print(\"Training took time \", datetime.now() - current_time)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CmbLTVniARy3",
        "colab_type": "text"
      },
      "source": [
        "Now let's use our test data to see how well our model did:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JIhejfpyJ8Bx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_input_fn = run_classifier.input_fn_builder(\n",
        "    features=test_features,\n",
        "    seq_length=MAX_SEQ_LENGTH,\n",
        "    is_training=False,\n",
        "    drop_remainder=False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PPVEXhNjYXC-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "estimator.evaluate(input_fn=test_input_fn, steps=None)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ueKsULteiz1B",
        "colab_type": "text"
      },
      "source": [
        "Now let's write code to make predictions on new sentences:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OsrbTD2EJTVl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def getPrediction(in_sentences):\n",
        "  labels = [\"Negative\", \"Positive\"]\n",
        "  input_examples = [run_classifier.InputExample(guid=\"\", text_a = x, text_b = None, label = 0) for x in in_sentences] # here, \"\" is just a dummy label\n",
        "  input_features = run_classifier.convert_examples_to_features(input_examples, label_list, MAX_SEQ_LENGTH, tokenizer)\n",
        "  predict_input_fn = run_classifier.input_fn_builder(features=input_features, seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=False)\n",
        "  predictions = estimator.predict(predict_input_fn)\n",
        "  return [(sentence, prediction['probabilities'], labels[prediction['labels']]) for sentence, prediction in zip(in_sentences, predictions)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-thbodgih_VJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pred_sentences = [\n",
        "  \"That movie was absolutely awful\",\n",
        "  \"The acting was a bit lacking\",\n",
        "  \"The film was creative and surprising\",\n",
        "  \"Absolutely fantastic!\"\n",
        "]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QrZmvZySKQTm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "predictions = getPrediction(pred_sentences)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MXkRiEBUqN3n",
        "colab_type": "text"
      },
      "source": [
        "Voila! We have a sentiment classifier!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ERkTE8-7oQLZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "predictions"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3Tg7c47vfLE",
        "colab_type": "text"
      },
      "source": [
        "# Export the model\n",
        "\n",
        "We are now ready to export the model. The following code defines the serving input function and exports the model to `OUTPUT_DIR`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NfXsdV4qtlpW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def serving_input_fn():\n",
        "    reciever_tensors = {\n",
        "        \"input_ids\": tf.placeholder(dtype=tf.int32,\n",
        "                                    shape=[1, MAX_SEQ_LENGTH])\n",
        "    }\n",
        "    features = {\n",
        "        \"input_ids\": reciever_tensors['input_ids'],\n",
        "        \"input_mask\": 1 - tf.cast(tf.equal(reciever_tensors['input_ids'], 0), dtype=tf.int32),\n",
        "        \"segment_ids\": tf.zeros(dtype=tf.int32, shape=[1, MAX_SEQ_LENGTH]),\n",
        "        \"label_ids\": tf.placeholder(tf.int32, [None], name='label_ids')\n",
        "    }\n",
        "    return tf.estimator.export.ServingInputReceiver(features, reciever_tensors)\n",
        "    \n",
        "estimator._export_to_tpu = False\n",
        "estimator.export_saved_model(OUTPUT_DIR+\"/export\", serving_input_fn)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tIFTmUbcwI0w",
        "colab_type": "text"
      },
      "source": [
        "# Upload the model to AWS\n",
        "\n",
        "Cortex loads models from AWS, so we need to upload the exported model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gByRzrnR_OBX",
        "colab_type": "text"
      },
      "source": [
        "Set these variables to configure your AWS credentials and model upload path:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1bdCOb3z0_Gh",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "AWS_ACCESS_KEY_ID = \"\" #@param {type:\"string\"}\n",
        "AWS_SECRET_ACCESS_KEY = \"\" #@param {type:\"string\"}\n",
        "S3_UPLOAD_PATH = \"s3://my-bucket/sentiment-analyzer/bert\" #@param {type:\"string\"}\n",
        "\n",
        "import sys\n",
        "import re\n",
        "\n",
        "if AWS_ACCESS_KEY_ID == \"\":\n",
        "    print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Please set AWS_ACCESS_KEY_ID\"), file=sys.stderr)\n",
        "\n",
        "elif AWS_SECRET_ACCESS_KEY == \"\":\n",
        "    print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Please set AWS_SECRET_ACCESS_KEY\"), file=sys.stderr)\n",
        "\n",
        "else:\n",
        "    try:\n",
        "        bucket, key = re.match(\"s3://(.+?)/(.+)\", S3_UPLOAD_PATH).groups()\n",
        "    except:\n",
        "        print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Invalid s3 path (should be of the form s3://my-bucket/path/to/file)\"), file=sys.stderr)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WLT09hZr_bhm",
        "colab_type": "text"
      },
      "source": [
        "Upload to S3:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jCN3BINl2sKN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import boto3\n",
        "\n",
        "s3 = boto3.client(\"s3\", aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY)\n",
        "\n",
        "for dirpath, _, filenames in os.walk(OUTPUT_DIR+\"/export\"):\n",
        "    for filename in filenames:\n",
        "        filepath = os.path.join(dirpath, filename)\n",
        "        filekey = os.path.join(key, filepath[len(OUTPUT_DIR+\"/export/\"):])\n",
        "        print(\"Uploading s3://{}/{} ...\".format(bucket, filekey), end = '')\n",
        "        s3.upload_file(filepath, bucket, filekey)\n",
        "        print(\" ✓\")\n",
        "\n",
        "print(\"\\nUploaded model export directory to \" + S3_UPLOAD_PATH)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7XPKSHzf_d7M",
        "colab_type": "text"
      },
      "source": [
        "<!-- CORTEX_VERSION_MINOR -->\n",
        "That's it! See the [example on GitHub](https://github.com/cortexlabs/cortex/tree/master/examples/tensorflow/sentiment-analyzer) for how to deploy the model as an API."
      ]
    }
  ]
}
